// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/exec/mark_distinct.h>
#include <pollux/common/base/range.h>
#include <pollux/vector/flat_vector.h>

#include <algorithm>
#include <utility>

namespace kumo::pollux::exec {
    MarkDistinct::MarkDistinct(
        int32_t operatorId,
        DriverCtx *driverCtx,
        const std::shared_ptr<const core::MarkDistinctNode> &planNode)
        : Operator(
            driverCtx,
            planNode->outputType(),
            operatorId,
            planNode->id(),
            "MarkDistinct") {
        const auto &inputType = planNode->sources()[0]->outputType();

        // Set all input columns as identity projection.
        for (auto i = 0; i < inputType->size(); ++i) {
            identityProjections_.emplace_back(i, i);
        }

        // We will use result[0] for distinct mask output.
        resultProjections_.emplace_back(0, inputType->size());

        groupingSet_ = GroupingSet::createForMarkDistinct(
            inputType,
            createVectorHashers(inputType, planNode->distinctKeys()),
            operatorCtx_.get(),
            &nonReclaimableSection_);

        results_.resize(1);
    }

    void MarkDistinct::addInput(RowVectorPtr input) {
        groupingSet_->addInput(input, false /*mayPushdown*/);

        input_ = std::move(input);
    }

    RowVectorPtr MarkDistinct::getOutput() {
        if (isFinished() || !input_) {
            return nullptr;
        }

        auto outputSize = input_->size();
        // Re-use memory for the ID vector if possible.
        VectorPtr &result = results_[0];
        if (result && result.use_count() == 1) {
            BaseVector::prepare_for_reuse(result, outputSize);
        } else {
            result = BaseVector::create(BOOLEAN(), outputSize, operatorCtx_->pool());
        }

        // newGroups contains the indices of distinct rows.
        // For each index in newGroups, we mark the index'th bit true in the result
        // vector.
        auto resultBits =
                results_[0]->as<FlatVector<bool> >()->mutableRawValues<uint64_t>();

        bits::fillBits(resultBits, 0, outputSize, false);
        for (const auto i: groupingSet_->hashLookup().newGroups) {
            bits::setBit(resultBits, i, true);
        }
        auto output = fillOutput(outputSize, nullptr);

        // Drop reference to input_ to make it singly-referenced at the producer and
        // allow for memory reuse.
        input_ = nullptr;

        return output;
    }

    bool MarkDistinct::isFinished() {
        return noMoreInput_ && !input_;
    }
} // namespace kumo::pollux::exec
